from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import PromptTemplate
from langchain_core.callbacks.stdout import StdOutCallbackHandler

_ = load_dotenv(find_dotenv())

prompt = PromptTemplate.from_template("请写一首关于{topic}的诗，并需要有标题。")

handlers = [StdOutCallbackHandler()]

chat = ChatZhipuAI(
    model="glm-4-plus",
    temperature=0.9,
    )

chain = prompt | chat
result = chain.invoke("春天", config={"callbacks": handlers})
print(result.content)